package com.base.saas.gateway.trace;

import org.springframework.cloud.client.loadbalancer.LoadBalanced;
import org.springframework.cloud.client.loadbalancer.reactive.LoadBalancedExchangeFilterFunction;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.reactive.function.client.ClientRequest;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurer;

/**
 * @author 王清亮
 * @date 2025/8/3 12:10
 * @description TODO
 */
@Configuration
public class TraceConfig {



    @Bean
    @LoadBalanced
    public WebClient.Builder loadBalancedWebClientBuilder(LoadBalancedExchangeFilterFunction lbFunction) {
        return WebClient.builder()
                .filter((request, next) -> {
                    // 添加 Trace 头
                    if (TraceContext.getTraceId() != null) {
                        ClientRequest filteredRequest = ClientRequest.from(request)
                                .headers(headers -> {
                                    headers.set(TraceContext.TRACE_ID, TraceContext.getTraceId());
                                    headers.set(TraceContext.SPAN_ID, TraceContext.getSpanId());
                                })
                                .build();
                        return next.exchange(filteredRequest);
                    }
                    return next.exchange(request);
                })
                .filter(lbFunction);
    }
}